#include <cuda_runtime.h>
#include <cuda_fp16.h> 
#include "conv2d.h"

extern "C" __global__ void implicit_gemm_v2(mykernelParamType param)
{
    const int ohow = blockIdx.x*16 + threadIdx.x;  // 输出空间位置(oh*ow)
    const int k = blockIdx.y*16 + threadIdx.y;     // 输出通道
    const int n = blockIdx.z;                      // batch索引
    
    const int tx = threadIdx.x;  // 线程的x坐标
    const int ty = threadIdx.y;  // 线程的y坐标
    __shared__ float sh_input[16][16];   // 16x16的输入数据缓存
    __shared__ float sh_weight[16][16];  // 16x16的权重数据缓存
    const int oh = ohow / param.Ow;  // 输出高度坐标
    const int ow = ohow % param.Ow;  // 输出宽度坐标
    
    const int ih_start = oh * param.u - param.p;  // 输入高度起始位置
    const int iw_start = ow * param.v - param.q;  // 输入宽度起始位置
    float sum = 0.0f;  // 累加器初始化

    const int input_offset = n * param.c * param.h * param.w;  // 当前batch的输入偏移
    const int kernel_offset = k * param.c * param.r * param.s; // 当前kernel的权重偏移
    const int crs_total = param.c * param.r * param.s;         // 总的内循环次数(c*r*s)

    for(int base=0; base<crs_total; base+=16){
        int crs_idx = base + tx;
        if(crs_idx < crs_total)
            sh_weight[ty][tx] = param.pweight[kernel_offset + crs_idx];  // 有效数据
        else 
            sh_weight[ty][tx] = 0.0f;  // 超出范围填充0
       int crs_idxy = base + ty;
        const int c = crs_idxy / (param.r * param.s);          // 输入通道
        const int r = (crs_idxy % (param.r * param.s)) / param.s;               
        const int s =crs_idxy % param.s;               
        const int ih = r + ih_start;  // 输入高度坐标
        const int iw = s+ iw_start;  // 输入宽度坐标
        if(ih>=0 && iw>=0 && ih<param.h && iw<param.w){
            const int i_addr =  c*param.h*param.w + ih*param.w + iw;
            sh_input[ty][tx] = param.pin[input_offset +i_addr];  // 有效输入数据
        } else {
            sh_input[ty][tx] = 0.0f;  // 超出边界填充0
        }
        __syncthreads();  // 确保所有线程完成共享内存写入
        #pragma unroll
        for(int i=0; i<16; i++){
            sum += sh_input[i][tx] * sh_weight[ty][i];
        }

        __syncthreads();  // 确保所有线程完成计算
    }
    if(ohow < param.Oh*param.Ow && k < param.k){
        const int o_addr = n * param.k * param.Oh * param.Ow +k * param.Oh * param.Ow +ohow;//oh * param.Ow + ow;                  // 空间位置偏移
        param.pout[o_addr] = sum;
    }
}

void launch_implicit_gemm_v2(unsigned int outh, unsigned int outw, unsigned int k, unsigned int n, mykernelParamType* param) {
    int blockx = (param->Oh * param->Ow + 15) / 16 ;  // 输出空间维度分块数
    int blocky = (param->k + 15) / 16;               // 输出通道维度分块数
    int blockz = param->n;                           // batch维度分块数
    int threadx = 16;  // x方向16线程
    int thready = 16;  // y方向16线程
    int threadz = 1;   // z方向1线程
    dim3 block(threadx, thready, threadz);  // 线程块维度
    dim3 grid(blockx, blocky, blockz);      // 网格维度
    implicit_gemm_v2<<<grid, block>>>(*param);
}